package com.liwy3.gateway.filter;


import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.cloud.gateway.filter.GatewayFilterChain;
import org.springframework.cloud.gateway.filter.GlobalFilter;
import org.springframework.core.Ordered;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferFactory;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
import org.springframework.http.HttpMethod;
import org.springframework.http.MediaType;
import org.springframework.http.server.reactive.ServerHttpRequestDecorator;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import org.springframework.stereotype.Component;
import org.springframework.web.server.ServerWebExchange;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.nio.CharBuffer;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.atomic.AtomicReference;

import static org.springframework.http.HttpHeaders.CONTENT_TYPE;


/**
 * Copyright © 元气森林（北京）食品科技集团有限公司. All rights reserved.
 *
 * @Description: 打印POST/JSON出入参, 注意postman请求header不能勾选 Accept-Encoding, 否则打印响应参数会乱码
 * @author: Crazy4J
 * @date: 2022年04月06日 10:40
 * @version: V1.0
 */
@Slf4j
@Component
public class BodyLoggingFilter implements GlobalFilter, Ordered {


    @Override
    public Mono<Void> filter(ServerWebExchange exchange, GatewayFilterChain chain) {


        AtomicReference<String> requestReference = new AtomicReference<>();

        ServerHttpRequestDecorator serverHttpRequestDecorator = new ServerHttpRequestDecorator(exchange.getRequest()) {
            @Override
            public Flux<DataBuffer> getBody() {
                return super.getBody().map(byteBuffer -> {
                    CharBuffer decode = StandardCharsets.UTF_8.decode(byteBuffer.asByteBuffer());
                    requestReference.set(decode.toString());
                    return byteBuffer;
                });
            }
        };

        ServerHttpResponseDecorator serverHttpResponseDecorator = new ServerHttpResponseDecorator(exchange.getResponse()) {
            @Override
            public Mono<Void> writeWith(@SuppressWarnings("NullableProblems") Publisher<? extends DataBuffer> body) {
                if (body instanceof Flux) {

                    DataBufferFactory bufferFactory = exchange.getResponse().bufferFactory();
                    String contentType = exchange.getRequest().getHeaders().getFirst(CONTENT_TYPE);
                    String method = exchange.getRequest().getMethodValue();

                    if (null != contentType
                            && HttpMethod.POST.name().equalsIgnoreCase(method)
                            && contentType.contains(MediaType.APPLICATION_JSON_VALUE)) {

                        Flux<? extends DataBuffer> fluxBody = Flux.from(body);
                        return super.writeWith(fluxBody.buffer().map(dataBuffers -> {

                            // 合并多个流集合，解决返回体分段传输
                            DataBufferFactory dataBufferFactory = new DefaultDataBufferFactory();
                            DataBuffer join = dataBufferFactory.join(dataBuffers);
                            byte[] content = new byte[join.readableByteCount()];
                            join.read(content);

                            // 释放掉内存
                            DataBufferUtils.release(join);

                            log.info("{},请求体:{}, 响应体:{}", exchange.getRequest().getURI().getRawPath(),
                                    requestReference.get(), new String(content, StandardCharsets.UTF_8));

                            return bufferFactory.wrap(content);
                        }));
                    }
                }
                return super.writeWith(body);
            }
        };

        return chain.filter(exchange.mutate().request(serverHttpRequestDecorator).response(serverHttpResponseDecorator).build());
    }


    @Override
    public int getOrder() {
        return HIGHEST_PRECEDENCE;
    }

}
